package com.bootx.predict;

import com.bootx.predict.pojo.PredictPlugin;
import com.bootx.predict.pojo.PredictionResult;
import com.bootx.predict.pojo.RedPacketBatch;
import com.bootx.predict.pojo.RedPacketRecord;
import com.bootx.predict.util.PredictUtils;
import org.springframework.stereotype.Component;

import java.util.*;

/**
 * @author black
 * 时间加权策略
 */
@Component("timeWeightedPredict")
public class TimeWeightedPredict extends PredictPlugin {

    private final double lambda = 0.001;

    @Override
    public String getName() {
        return "timeWeightedPredict";
    }

    @Override
    public List<PredictionResult> predict(List<RedPacketBatch> list, Set<Integer> indexes) {
        List<RedPacketRecord> history = PredictUtils.parseData(list);
        // 获取最大 openTime，用于时间衰减计算
        int maxTime = history.stream().mapToInt(RedPacketRecord::getOpenTime).max().orElse(0);

        // 数据结构：index -> {weightSum, oddWeightedSum, openTimes[], total, oddCount}
        Map<Integer, Double> weightedOddSum = new HashMap<>();
        Map<Integer, Double> weightedTotalSum = new HashMap<>();
        Map<Integer, List<Integer>> timeMap = new HashMap<>();
        Map<Integer, Integer> totalCount = new HashMap<>();
        Map<Integer, Integer> oddCountMap = new HashMap<>();

        for (RedPacketRecord r : history) {
            double weight = Math.exp(-lambda * (maxTime - r.getOpenTime()));
            double oddValue = PredictUtils.isAmountOdd(r.getAmount()) ? 1.0 : 0.0;

            weightedOddSum.merge(r.getIndex(), weight * oddValue, Double::sum);
            weightedTotalSum.merge(r.getIndex(), weight, Double::sum);

            timeMap.computeIfAbsent(r.getIndex(), k -> new ArrayList<>()).add(r.getOpenTime());
            totalCount.merge(r.getIndex(), 1, Integer::sum);
            if (oddValue == 1.0) {
                oddCountMap.merge(r.getIndex(), 1, Integer::sum);
            }
        }

        List<PredictionResult> results = new ArrayList<>();
        for (int idx : indexes) {
            double num = weightedOddSum.getOrDefault(idx, 0.0);
            double den = weightedTotalSum.getOrDefault(idx, 0.0);
            double probability = den > 0 ? num / den : 0.5;
            List<Integer> times = timeMap.getOrDefault(idx, Collections.emptyList());
            double avgTime = times.stream().mapToInt(Integer::intValue).average().orElse(0);
            int total = totalCount.getOrDefault(idx, 0);
            int oddCount = oddCountMap.getOrDefault(idx, 0);
            results.add(new PredictionResult(idx, total, oddCount, probability, Double.valueOf(avgTime+"").intValue()));
        }

        return results;
    }
}
